import cv2
import numpy as np
import pickle

# 寻找最大轮廓
def FindBigestContour(src):
    imax = 0
    imaxcontour = -1

    contours, hierarchy = cv2.findContours(src, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    for i in range(len(contours)):
        itmp = cv2.contourArea(contours[i])
        if imaxcontour < itmp:
            imax = i
            imaxcontour = itmp
    print(imax)
    return contours[imax]

# 去光差
def moveLightDiff(src):
    srcclone = src.copy()

    kernel = np.ones((100, 100), np.uint8)  # 用于腐蚀计算的核
    srcclone = cv2.erode(srcclone,  # 原始图像
                        kernel,  # 腐蚀核
                        iterations=5)  # 迭代次数
    srcclone = cv2.dilate(srcclone,  # 原始图像
                         kernel,  # 腐蚀核
                         iterations=5)  # 迭代次数
    cv2.imshow('srcclone', srcclone)
    cv2.waitKey(0)
    dst = src - srcclone
    return dst

"""
图片始终为黑色背景
"""
def blackBackground(src, src_h):
    height = src.shape[0]  # 图片大小
    width = src.shape[1]

    std, stddev = cv2.meanStdDev(src_h)  # 均值std
    b = std[0]

    for row in range(height):
        for col in range(width):
            tmp = src_h[row, col]
            src_h[row, col] = src_h[row, col] - b
            if src_h[row, col] < 0:
                src_h[row, col] = tmp
    # print("src_h:", src_h)
    return src_h

"""
计算西瓜颜色均值
"""
def green(src):
    hsv = cv2.cvtColor(src, cv2.COLOR_BGR2HSV)
    lower = np.array([0,43,46])
    higher = np.array([155,255,255])
    mask = cv2.inRange(hsv, lowerb=lower, upperb=higher)
    cv2.imshow('green', mask)
    cv2.waitKey(0)

    # 计算颜色均值
    mean = cv2.mean(src, mask=mask)
    print("average: ", mean)
    return 1

"""
判断纹理是否清晰
"""
def grain(src):
    # 灰度化和LBP特征提取
    gray = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)
    radius = 1
    n_points = 8 * radius
    lbp = cv2.ORB_create(n_points, radius)
    keypoints = lbp.detect(gray, None)
    mask = cv2.drawKeypoints(gray, keypoints, None, color=(0, 255, 0), flags=0)
    cv2.imshow("grain", mask)
    cv2.waitKey(0)

    # 计算纹理清晰度
    mean = cv2.mean(mask)[0]

    # 输出结果
    if mean > 00:
        print(1)
        return 1
    else:
        print(2)
        return 2

"""
判断是圆形还是椭圆
"""
def circle(cnt):
    # 使用cv2.fitEllipse()函数拟合椭圆，计算椭圆的长轴和短轴长度，比较大小关系
    ellipse = cv2.fitEllipse(cnt)
    (x, y), (a, b), angle = ellipse
    print("距离：", abs(a-b))
    if abs(a-b) < 60:
        # print(1)
        return 1
    else:
        # print(2)
        return 2

def watermelon_main(img):
    src = cv2.imread(img)
    gray = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)
    # blur = cv2.GaussianBlur(gray, (3, 3), 0)    # 高斯模糊
    # canny = cv2.Canny(blur, 50, 150)   # 边缘检测
    # cv2.imshow("blur", blur)
    # cv2.waitKey(0)

    src_hsv = cv2.cvtColor(src, cv2.COLOR_BGR2HSV)

    # h通道图片
    rgb_planes = cv2.split(src_hsv)
    src_h = rgb_planes[0]

    # 去光差
    # src_h = moveLightDiff(src_h)
    # cv2.imshow('after_move', src_h)
    # cv2.waitKey(0)

    # 背景始终为黑色，目标始终为白色
    # src_h = blackBackground(src, src_h)
    # cv2.imshow('black:', src_h)
    # cv2.waitKey(0)

    # # 定义结构元素
    # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (20, 20))
    # # 开闭运算，先开运算去除背景噪声，再继续闭运算填充目标内的孔洞
    # opened = cv2.morphologyEx(src_h, cv2.MORPH_OPEN, kernel)
    # src_h = cv2.morphologyEx(opened, cv2.MORPH_CLOSE, kernel)
    # cv2.imshow("closed", src_h)
    # cv2.waitKey(0)

    # 二值化
    ret, thresh = cv2.threshold(src_h, 100, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    # 锐化处理
    thresh = cv2.subtract(thresh, gray)

    # 腐蚀和膨胀
    kernel = np.ones((7, 7), np.uint8)  # 设置kernel大小
    thresh = cv2.erode(thresh, kernel, iterations=3)  # 腐蚀
    thresh = cv2.dilate(thresh, kernel, iterations=1)  # 膨胀还原图形

    # 最大轮廓
    bigestcontrour = FindBigestContour(thresh)
    # contours, hierarchy = cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_NONE)

    srccopy = src.copy()
    # # 近似轮廓
    # epsilon = 0.002 * cv2.arcLength(bigestcontrour, True)
    # bigestcontrour = cv2.approxPolyDP(bigestcontrour, epsilon, True)

    # 绘制轮廓
    cv2.drawContours(srccopy, [bigestcontrour], -1, (0, 0, 255), 3)
    cv2.imshow('1', srccopy)
    cv2.waitKey(0)

    mask = np.zeros_like(src)  # 用numpy生成一个全为0的，300*300的二维数组
    cv2.drawContours(mask, [bigestcontrour], -1, (255, 255, 255), 2)    # 根据绘制出的轮廓创建掩膜
    cv2.fillPoly(mask, [bigestcontrour], (255, 255, 255))     # 掩膜轮廓内部填充为白色
    img1 = cv2.bitwise_and(src, mask)  # 将图像与掩模经行按位与操作
    cv2.imshow('img1', img1)
    cv2.waitKey(0)

    greencolor = green(img1)    # 获取到bgr三通道的均值
    clear = grain(img1)   # 用lbp（局部二值模式）进行纹理提取
    form = circle(bigestcontrour)
    print("color=",greencolor,"清晰度=",clear,"形状=",form)

    cv2.destroyAllWindows()



    # 加载训练好的模型
    with open('decision_tree_model.pkl', 'rb') as f:
        clf = pickle.load(f)

    # 输入待预测的数据
    x1 = greencolor
    x2 = clear
    x3 = form
    X_test = [[x1, x2, x3]]

    # 使用模型进行预测
    y_pred = clf.predict(X_test)
    print(y_pred)

    return y_pred[0]
